package internal

import (
	"github.com/gin-gonic/gin"
	"net/http"
	"time"
)

// RateLimitMiddleware 限流中间件
func RateLimitMiddleware(r *gin.Engine, rate int, capacity int) {
	// 创建令牌桶
	limiter := NewLimiter(rate, capacity)

	// 使用中间件
	r.Use(func(c *gin.Context) {
		// 从桶中取出一个令牌
		if !limiter.Allow() {
			// 如果桶中没有令牌，则返回错误
			c.JSON(http.StatusTooManyRequests, gin.H{"msg": "Too many requests", "type": "rate"})
			c.Abort()
			return
		}

		// 正常处理请求
		c.Next()
	})
}

// Limiter 令牌桶限流器
type Limiter struct {
	capacity int
	tokens   int
	rate     time.Duration
	next     int64
}

// NewLimiter 创建一个新的限流器
func NewLimiter(rate int, capacity int) *Limiter {
	return &Limiter{
		capacity: capacity,
		tokens:   capacity,
		rate:     time.Duration(rate) * time.Second,
		next:     time.Now().Unix(),
	}
}

// Allow 检查是否可以从桶中取出一个令牌
func (l *Limiter) Allow() bool {
	now := time.Now().Unix()
	if now > l.next {
		// 重置桶中的令牌数量
		l.tokens = l.capacity
		l.next = now + int64(l.rate.Seconds())
	}

	if l.tokens == 0 {
		// 桶中没有令牌
		return false
	}

	// 从桶中取出一个令牌
	l.tokens--
	return true
}
